import matplotlib.pyplot as plt
import pandas as pd

class EnhancedDataVisualizer:
    def __init__(self, excel_name, sheet_name, li, chart_type='bar', title_mapping=None):
        self.excel_name = excel_name
        self.sheet_name = sheet_name
        self.li = li if isinstance(li, list) else [li]
        self.chart_type = chart_type.lower()
        self.title_mapping = title_mapping or {}
        self.validate_chart_type()

    def validate_chart_type(self):
        valid_types = ['bar', 'pie']
        if self.chart_type not in valid_types:
            raise ValueError(f"无效图表类型，可选值：{', '.join(valid_types)}")

    def get_df(self):
        df = pd.read_excel(self.excel_name, sheet_name=self.sheet_name)
        for col in self.li:
            if pd.api.types.is_string_dtype(df[col]):
                try:
                    df[col] = pd.to_numeric(df[col], errors='ignore')
                except:
                    pass
        return df

    def get_data(self, df):
        if self.chart_type == 'pie':
            return df[self.li[0]].value_counts()
        else:
            return df[self.li].sum()

    def config_base_style(self):
        plt.rcParams['font.family'] = 'Microsoft YaHei'
        plt.figure(figsize=(10, 7), facecolor='#F5F5F5')  # 调整画布比例

    def generate_bar(self, data):
        plt.bar(data.index, data.values,
                color='#4B9DA3', edgecolor='gray',
                linewidth=1.2, alpha=0.8)
        for x, y in enumerate(data.values):
            plt.text(x, y, f'{int(y)}', ha='center', va='bottom', fontsize=12)

    def generate_pie(self, data):
        wedges, texts, autotexts = plt.pie(
            data.values,
            startangle=90,
            colors=plt.cm.Pastel1.colors,
            wedgeprops={'linewidth': 1, 'edgecolor': 'gray'},
            autopct=lambda p: f'{p:.1f}%({int(p * sum(data.values) / 100)})',
            textprops={
                'fontsize': 5,  # 字体大小调整为5pt
                'color': 'black'  # 增加字体颜色定义
            }
        )

        plt.legend(
            wedges,
            data.index,
            title=self.title_mapping.get(self.li[0], self.li[0]),
            loc='center left',
            bbox_to_anchor=(1, 0.5),
            frameon=False,
            fontsize=8  # 同步调整图例字体大小
        )
        plt.axis('equal')

    def save_chart(self):  # 修复缩进
        suffix = self.chart_type
        plt.savefig(f'{self.li[0]}_{suffix}.png', dpi=300, bbox_inches='tight')

    def visualize(self):
        df = self.get_df()
        processed_data = self.get_data(df)

        self.config_base_style()
        if self.chart_type == 'bar':
            self.generate_bar(processed_data)
            title_suffix = "数值总和分布"
        else:
            self.generate_pie(processed_data)
            title_suffix = "分类占比分布"

        # 动态标题配置
        display_name = self.title_mapping.get(self.li[0], self.li[0])
        plt.title(f"{display_name} {title_suffix}", fontsize=14, pad=20)
        plt.tight_layout()
        self.save_chart()
        plt.show()

if __name__ == '__main__':
    # 数值型列配置
    number_groups = [
        ['S1', 'S2', 'S3', 'S4', 'S5', 'S6', 'S7'],
        ['C1', 'C2', 'C3', 'C4', 'C5', 'C6', 'C7'],
        ['P1', 'P2', 'P3', 'P4', 'P5', 'P6', 'P7', 'P8'],
        ['A1', 'A2', 'A3', 'A4', 'A5', 'A6', 'A7', 'A8', 'A9'],
        ['Q1', 'Q2', 'Q3', 'Q4', 'Q5', 'Q6']
    ]

    # 分类列标题映射
    category_titles = {
        'sex': '性别分布',
        'age': '年龄分层',
        'career': '职业构成'
    }

    # 生成数值型柱状图
    for group in number_groups:
        visualizer = EnhancedDataVisualizer('clear_maded_food.xlsx', '预制表格', group)
        visualizer.visualize()

    # 生成分类饼状图（带右侧图例）
    for column in category_titles.keys():
        visualizer = EnhancedDataVisualizer(
            excel_name='clear_maded_food.xlsx',
            sheet_name='预制表格',
            li=column,
            chart_type='pie',
            title_mapping=category_titles
        )
        visualizer.visualize()